import torch
import cv2
import os

# 加载 YOLOv5 模型
model = torch.hub.load(
    r"C:\Users\23326\Desktop\yolov5-7.0", 
    'custom', 
    path=r"C:\Users\23326\Desktop\yolov5-7.0\runs\train\exp10\weights\best.pt", 
    source='local'
)


# 设置置信度阈值
model.conf = 0.4

# 原图路径（绝对路径）
img_path = r"C:\Users\23326\Desktop\yolov5-7.0\dataset\images\test\0125-91_89-321-300_510-376-510-372_323-368_327-305_514-309-10_8_6_18_25_33_25-95-46_jpg.rf.184d65fe2b4fad079c4836c36f164b42.jpg"
img = cv2.imread(img_path)

# 推理
results = model(img)

# 获取预测结果
detections = results.xyxy[0]  # tensor格式: [x1, y1, x2, y2, conf, class]

# 创建保存目录（绝对路径）
save_dir = r"C:\Users\23326\Desktop\yolov5-7.0\cropped_plates"
os.makedirs(save_dir, exist_ok=True)

# 遍历检测结果
for i, det in enumerate(detections):
    x1, y1, x2, y2 = map(int, det[:4])
    cropped = img[y1:y2, x1:x2]

    # 保存裁剪图像
    save_path = os.path.join(save_dir, f"plate_{i}.jpg")
    cv2.imwrite(save_path, cropped)
    print(f"已保存裁剪车牌图像：{save_path}")


from paddleocr import PaddleOCR
import cv2
import os

# 初始化 OCR 模型（中文，支持旋转检测）
ocr = PaddleOCR(use_angle_cls=True, lang='ch')

# 设置裁剪图像目录
crop_dir = r"C:\Users\23326\Desktop\yolov5-7.0\cropped_plates"

# 遍历所有裁剪图像
for img_name in os.listdir(crop_dir):
    img_path = os.path.join(crop_dir, img_name)
    img = cv2.imread(img_path)

    if img is None:
        print(f"图像加载失败: {img_path}")
        continue

    # 转为 RGB
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # OCR 识别
    result = ocr.ocr(img_rgb, cls=True)

    print(f"【{img_name}】识别结果：")
    for line in result:
        for word_info in line:
            text, score = word_info[1]
            print(f"  文本: {text}，置信度: {score:.2f}")
